import numpy as np
import os
from cloud_detection import *
from sun_position_identification import *
from vis import vis_single_img, classify_weather, classify_weather_2
from collections import Counter
import csv
import h5py

# 配置路径
def configure_paths():
    cwd = os.getcwd()
    pardir = os.path.dirname(cwd)
    data_dir = os.path.join(pardir, "data", "video_data")
    figs_dir = os.path.join(pardir, "figs")
    
    paths = {
        'test': {
            'time': os.path.join(data_dir, 'times_5min_test.npy'),
            'img': os.path.join(data_dir, 'predicted_images.npy'),  # 假设测试集图像文件名为 predicted_images_test.npy
        },
        'train': {
            'time': os.path.join(data_dir, 'times_5min_trainval.npy'),
            'img': None  # 这里先设置为 None，后续用 images_pred_train 替代
        }
    }
    return paths, figs_dir

# 加载数据集
def load_dataset(path_dict, images_pred_train=None):
    data = {}
    for key in path_dict:
        if key == 'train' and images_pred_train is not None:
            # 对于 train 数据集，使用预加载的 images_pred_train
            data[key] = {
                'time': np.load(path_dict[key]['time'], allow_pickle=True),
                'images': images_pred_train
            }
        else:
            # 对于 test 数据集，从 .npy 文件加载图像数据
            data[key] = {
                'time': np.load(path_dict[key]['time'], allow_pickle=True),
                'images': np.load(path_dict[key]['img'], allow_pickle=True)
            }
        print(f"{key} time stamps shape: {data[key]['time'].shape}")
        print(f"{key} images shape: {data[key]['images'].shape}")
    return data

# 处理所有示例并保存结果
def show_all_examples(data, batch_window, dataset_name):
    all_cloud_covers = []  # 用于存储所有的 cloud_cover
    all_timestamps = []  # 用于存储所有的 timestamp
    all_weathers = []
    
    for key in data:
        # 处理第一个 batch 的前 3 个窗口的图像
        images_first_window = data[key]['images'][0, :batch_window - 1, :, :, :]  # 取第一个 batch 的前 3 个窗口的图像
        timestamp_first_window = data[key]['time'][0]  # 使用第一个时间戳
        print(f"Displaying first 3 windows of {key} samples from the first batch:")
        for i, img in enumerate(images_first_window):
            cloud_cover, SMAPI = vis_single_img(timestamp_first_window, key, img)  # 使用同一个时间戳
            print("cloud_cover", cloud_cover)
            all_weathers.append(classify_weather(cloud_cover, SMAPI))
            all_cloud_covers.append(cloud_cover)  # 收集 cloud_cover
        
        # 处理所有 batch 的最后一个窗口的图像
        images_last_window = data[key]['images'][:, -1, :, :, :]
        print("images_last_window", images_last_window.shape)
        times = data[key]['time']
        for i, (img, ts) in enumerate(zip(images_last_window, times)):
            cloud_cover, SMAPI = vis_single_img(ts, key, img)
            print("SMAPI", SMAPI)
            all_weathers.append(classify_weather(cloud_cover, SMAPI))
            all_cloud_covers.append(cloud_cover)  # 收集 cloud_cover
            all_timestamps.append(ts)  # 收集时间戳

    sliding_windows = []
    for i in range(len(all_weathers) - batch_window + 1): 
        window = all_weathers[i:i + batch_window]
        sliding_windows.append({
            'cloud_covers': window,
            'timestamp': all_timestamps[i]
        })
    
    for window_data in sliding_windows:
        window = window_data['cloud_covers']
        weather_counts = Counter(window)
        most_common_weather, _ = weather_counts.most_common(1)[0]  # 获取最常见的天气类型
        window_data['dominant_weather'] = most_common_weather
        del window_data["cloud_covers"]
    
    # 输出结果
    for idx, window_data in enumerate(sliding_windows):
        if idx % 10 == 0:
            print(f"Window {idx}: Timestamp={window_data['timestamp']}, Dominant Weather={window_data['dominant_weather']}")

    # 定义文件名
    csv_file = f"weather_data_{dataset_name}.csv"

    # 写入 CSV 文件
    with open(csv_file, mode='w', newline='', encoding='utf-8') as file:
        fieldnames = ['timestamp', 'dominant_weather']  # 列名
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        
        # 写入表头
        writer.writeheader()
        
        # 写入数据
        for data in sliding_windows:
            writer.writerow(data)

    print(f"数据已成功保存到 {csv_file}")
    print(len(sliding_windows))

if __name__ == "__main__":
    # 初始化配置
    batch_window = 3
    paths, output_dir = configure_paths()

    # 加载 HDF5 文件以获取 images_pred_train
    data_folder = os.path.join(os.path.dirname(os.path.dirname(os.getcwd())), 'solar_former', 'data', 'video_data')
    data_path = os.path.join(data_folder, 'video_5min.hdf5')
    with h5py.File(data_path, 'r') as f:
        trainval = f['trainval']
        images_pred_train = trainval['image_pred_trainval'][...][:, ::2, :, :, :]  # 提取训练集预测图像

    # 加载数据集
    dataset = load_dataset(paths, images_pred_train=images_pred_train)

    # 先处理 test 数据
    print("处理 test 数据...")
    show_all_examples({key: dataset[key] for key in ['test']}, batch_window, "test")

    # 再处理 train 数据
    print("处理 train 数据...")
    show_all_examples({key: dataset[key] for key in ['train']}, batch_window, "train")
    

